from sklearn.datasets import make_moons
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt


def load_dataset(samples=2000, ratio=0.2):
    """
    加载 sklearn 中的 make_moons 数据集,按照比例进行分割
    :param samples: 数据的采样个数
    :param ratio: 测试数据所占比例
    :return: 数据集所有数据，数据集所有数据对应的标签，训练集数据，训练集数据标签，测试集数据，测试集数据标签
    """
    samples = 2000
    X, y = make_moons(n_samples=samples, noise=0.2, random_state=100)
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=ratio, random_state=42)
    # print(X_train.shape)
    # print(X_test.shape)
    # print(y_train.shape)
    # print(y_test.shape)

    return X, y, X_train, X_test, y_train, y_test


def make_plot(X, y, plot_name, XX=None, YY=None, preds=None, dark_mode=False):
    """
    将数据绘制为散点图
    :param X:
    :param y:
    :param plot_name:
    :param XX:
    :param YY:
    :param preds:
    :param dark_mode: 是否为黑暗模式
    :return:
    """
    # 查看 matplotlib 支持的风格
    # print(plt.style.available)
    if (dark_mode):
        plt.style.use("dark_background")
    else:
        # sns.set_style("whitegrid")
        plt.style.use("bmh")

    # 设置图的大小
    plt.figure(figsize=(12, 7))
    axes = plt.gca()
    axes.set(xlabel="$x_1$", ylabel="$x_2$")
    plt.title(plot_name, fontsize=30)
    # 绘图内容从左边 10% 开始
    plt.subplots_adjust(left=0.10)
    # 绘图内容到右边 10% 结束
    plt.subplots_adjust(right=0.90)

    if XX is not None and YY is not None and preds is not None:
        plt.contourf(XX, YY, preds.reshape(XX.shape), 25, alpha=1, cmap=plt.cm.Spectral)
        plt.contour(XX, YY, preds.reshape(XX.shape), levels=[.5], cmap="Greys", vmin=0, vmax=.6)
    # 绘制散点图，根据标签区分颜色
    plt.scatter(X[:, 0], X[:, 1], c=y.ravel(), s=40, cmap=plt.cm.Spectral, edgecolors='none')
    plt.show()


if __name__ == '__main__':
    X, y, X_train, X_test, y_train, y_test = load_dataset()
    # print(X.shape)
    # print(y.shape)
    # print(X_train.shape)
    # print(X_test.shape)
    # print(y_train.shape)
    # print(y_test.shape)
    # print(y_train)

    make_plot(X, y, "Classification Dataset Visualization")
